import streamlit as st
import pickle
from PIL import Image
import numpy as np
from io import BytesIO

# 加载模型
with open('Downloads/Data_Rate/best_model.pkl', 'rb') as file:
    rfr_model = pickle.load(file)

# 定义计算颜色直方图的函数
def compute_color_histogram_pil(image, bins=32):
    # Convert the image to the RGB color space
    r, g, b = image.split()
    # Compute the color histogram
    hist_r = r.histogram()
    hist_g = g.histogram()
    hist_b = b.histogram()
    # Normalize the histograms
    hist_r = [x / sum(hist_r) for x in hist_r]
    hist_g = [x / sum(hist_g) for x in hist_g]
    hist_b = [x / sum(hist_b) for x in hist_b]
    # Combine the histograms
    hist_combined = np.concatenate((hist_r, hist_g, hist_b))
    return hist_combined

# Streamlit界面
st.title('Image Attractiveness Predictor')

# 文件上传控件
uploaded_file = st.file_uploader("Please upload an image", type=["png", "jpg", "jpeg"])

if uploaded_file is not None:
    # 显示上传的图片
    image = Image.open(uploaded_file)
    st.image(image, caption='Uploaded Image', use_column_width=True)

    # 计算新图像的颜色直方图
    new_image_histogram = compute_color_histogram_pil(image)

    # 归一化直方图
    new_image_histogram_normalized = new_image_histogram / np.linalg.norm(new_image_histogram)

    # 将直方图转换为适合模型输入的格式
    new_image_features = new_image_histogram_normalized.reshape(1, -1)

    # 使用处理后的输入数据进行预测
    predicted_attractiveness = rfr_model.predict(new_image_features)
    predicted_score = predicted_attractiveness[0]

    # 假设原始分数范围是0到5
    min_score = 0
    max_score = 5

    # 转换为百分制
    percentage_score = round((predicted_score - min_score) / (max_score - min_score) * 100+30, 2)

    # 显示预测结果
    st.write(f"Predicted Attractiveness Score: {percentage_score:.2f}")

